﻿#include "pch.h"
#include "TestComExport.h"
#include "CTestCom.h"
#include "CTestComFactory.h"

/**
* DLL操作：
	注册DLL：regsvr32.exe TestCom_Server.dll
	注销DLL：regsvr32.exe TestCom_Server.dll -u
**/

/**
* CLSID项下子项意义：
	InprocServer32				如果这是一个 DLL 类型的 COM 对象，表示其路径
	LocalServer32				如果这是一个 EXE 类型的 COM 对象，表示其路径
	ProgID						映射到的 ProgID，例如 Excel.Application
	VersionIndependentProgID	不带版本号的 ProgID
	TypeLib						指定 Type Library（接口定义的描述库）
	Implemented Categories		支持的 COM 分类（用于组件管理）
**/

// DLL进程实例句柄
static HMODULE g_hModule;
// 组件中CTestCom对象个数
static ULONG g_num;

/******************** 注册表辅助函数 ***********************/

static bool CheckRegistryWritePermission(HKEY root, LPCWSTR subkey) {
    HKEY hKey;
    LONG lRes = RegOpenKeyExW(root, subkey, 0, KEY_WRITE, &hKey);
    if (lRes == ERROR_SUCCESS) {
        RegCloseKey(hKey);
        return true;
    } else {
        std::wcout << L"[权限不足] 无法写入注册表: " << subkey << L" 错误码: " << lRes << std::endl;
        return false;
    }
}

static void ComInfoRegRegister(LPCWSTR lpPath) {
    HKEY thk, tclsidk;
    bstr_t desc(L"Test Com Component Object");
    // 注册CLSID信息
    if (CheckRegistryWritePermission(HKEY_CLASSES_ROOT, L"CLSID")) {
        if (RegOpenKeyExW(HKEY_CLASSES_ROOT, L"CLSID", 0, KEY_WRITE, &thk) == ERROR_SUCCESS) {
            if (RegCreateKeyExW(thk, CLSID_TESTCOM_STR, 0, NULL, 0, KEY_WRITE, NULL, &tclsidk, NULL) == ERROR_SUCCESS) {
                RegSetValueExW(tclsidk, NULL, 0, REG_SZ, (const BYTE*)desc.GetBSTR(), (DWORD)(wcslen(desc) + 1) * 2);
                HKEY tinps32k;
                if (RegCreateKeyExW(tclsidk, L"InprocServer32", 0, NULL, 0, KEY_WRITE, NULL, &tinps32k, NULL) == ERROR_SUCCESS) {
                    RegSetValueExW(tinps32k, NULL, 0, REG_SZ, (const BYTE*)lpPath, (DWORD)(wcslen(lpPath) + 1) * 2);
                    std::wcout << L"Com Infomation Register Success!" << std::endl;
                    RegCloseKey(tinps32k);
                }
                RegCloseKey(tclsidk);
            }
            RegCloseKey(thk);
        }
    }
    // 注册ProgID信息
    if (CheckRegistryWritePermission(HKEY_CLASSES_ROOT, NULL)) {
        if (RegCreateKeyExW(HKEY_CLASSES_ROOT, L"COMCTL.CTestCom", 0, NULL, 0, KEY_WRITE, NULL, &thk, NULL) == ERROR_SUCCESS) {
            RegSetValueExW(thk, NULL, 0, REG_SZ, (const BYTE*)desc.GetBSTR(), (DWORD)(wcslen(desc) + 1) * 2);
            if (RegCreateKeyExW(thk, L"CLSID", 0, NULL, 0, KEY_WRITE, NULL, &tclsidk, NULL) == ERROR_SUCCESS) {
                RegSetValueExW(tclsidk, NULL, 0, REG_SZ, (const BYTE*)CLSID_TESTCOM_STR, (DWORD)(wcslen(CLSID_TESTCOM_STR) + 1) * 2);
                std::wcout << L"Com ProgID Register Success!" << std::endl;
                RegCloseKey(tclsidk);
            }
            RegCloseKey(thk);
        }
    }
}

static void _hKeyDelete(HKEY hk, LPCWSTR lp) {
    RegDeleteKeyW(hk, lp);
}

static void ComInfoRegUnRegister() {
    HKEY thk;
    // 删除CLSID信息
    if (RegOpenKeyExW(HKEY_CLASSES_ROOT, L"CLSID", 0, KEY_WRITE, &thk) == ERROR_SUCCESS) {
        bstr_t bstr(CLSID_TESTCOM_STR);
        bstr += L"\\InprocServer32";
        _hKeyDelete(thk, bstr);
        _hKeyDelete(thk, CLSID_TESTCOM_STR);
        RegCloseKey(thk);
    }
    // 删除ProgID信息
    if (RegOpenKeyExW(HKEY_CLASSES_ROOT, L"COMCTL.CTestCom", 0, KEY_WRITE, &thk) == ERROR_SUCCESS) {
        _hKeyDelete(thk, L"CLSID");
        RegCloseKey(thk);
    }
    _hKeyDelete(HKEY_CLASSES_ROOT, L"COMCTL.CTestCom");
}

/********************** 注册函数 **************************/

HRESULT __stdcall DllRegisterServer()
{
    WCHAR szModule[1024];
    DWORD dwResult = GetModuleFileNameW(g_hModule, szModule, 1024);
    if (dwResult == 0) {
        return -1;
    }
    MessageBoxW(NULL, szModule, L"COM Register", MB_OK);
    ComInfoRegRegister(szModule);
    return 0;
}

HRESULT __stdcall DllUnregisterServer()
{
	ComInfoRegUnRegister();
	return 0;
}

HRESULT __stdcall DllCanUnloadNow()
{
	if (g_num == 0)
		return S_OK;
	else
		return S_FALSE;
}

HRESULT __stdcall DllGetClassObject(REFCLSID rclsid, REFIID riid, LPVOID* ppv)
{
	LPOLESTR szCLSID;
	(void)StringFromCLSID(rclsid, &szCLSID);
	wprintf_s(L"Get RCLSID is \"%s\" \n", szCLSID);

	if (rclsid == CLSID_TESTCOM) {
		CTestComFactory* factory = new CTestComFactory();
		if (factory == NULL) {
			return E_OUTOFMEMORY;
		}
		HRESULT hres = factory->QueryInterface(riid, ppv);
		return hres;
	}
	else {
		return CLASS_E_CLASSNOTAVAILABLE;
	}
}

// DLL入口点
BOOL APIENTRY DllMain(HMODULE hModule,
	DWORD  ul_reason_for_call,
	LPVOID lpReserved
)
{
	// 设置DLL进程实例句柄
	g_hModule = hModule;

	switch (ul_reason_for_call)
	{
	case DLL_PROCESS_ATTACH:
	case DLL_THREAD_ATTACH:
	case DLL_THREAD_DETACH:
	case DLL_PROCESS_DETACH:
		break;
	}
	return TRUE;
}
